from virtual_system import System, analysis

from dwave.system.composites import EmbeddingComposite
# from dwave.system.samplers import DWaveSampler
import neal
import dwavebinarycsp
from penalty_function import makeExactOne, gridSweep
from collections import defaultdict
import dimod
import random


sampler = neal.SimulatedAnnealingSampler()
# node读取时间、remote node节点的时间
def getNodeQubit(own_node, read_node, time_frame):
    return f'Node{own_node}_Read{read_node}_Time{time_frame}'


def praiseQubitName(qubit_name):
    own_node, read_node, time = qubit_name.split('_')

    return int(own_node.replace('Node', '')), int(read_node.replace('Read', '')), int(time.replace('Time', ''))

def praise2frame_qubit(sample, node_number, time_frame_number):
    frame2node2assign_node = [
        [
            None
            for _ in range(node_number)
        ]
        for _ in range(time_frame_number)
    ]
    # 麻了，乘号里面竟然是引用复制

    for qubit_name, value in sample.items():
        own_node, read_node, time = praiseQubitName(qubit_name)

        # frame2node2read_node[time][own_node][read_node] = value

        if value == 1:
            assert frame2node2assign_node[time][own_node] == None
            frame2node2assign_node[time][own_node] = read_node
    

    return frame2node2assign_node

# 先是不考虑资源的
class QuantumSystemNoResource(System):
    former_hearbeats = None


    fixed_qubits = {}

    def refreshBQM(self):
        self.fixed_qubits = {}

        self.inner_frame_index = 0  #相较于上一次的心跳的时间
        heartbeats = self.heartbeats

        TIME_FRAME_NUMBER = self.heartbeat_rate  # 假设 0.1 us 一段
        TIME_FRAMES = [i for i in range(TIME_FRAME_NUMBER)]

        NODE_NUMBER = len(self.nodes)
        NODES = [i for i in range(NODE_NUMBER)]

        QUEUE_SIZE_PER_NODE = {
            node.index: heartbeats[node.index]['queue_size']
            for node in self.nodes
        }

        AVG_EXECUTION_TIME_PER_BODE = {
            node.index: heartbeats[node.index]['avg_execution_time']
            for node in self.nodes 
        }

        # 初始化网速
        NETWORK_SPEED = self.network_speed


        # 每个节点的queue的数量 * 平均时间 要小
        # 每个节点的利用率高（1的数量多）

        # 每次都的时候要考虑 （传输速度，等待时间）,等待时常是考虑了之前的read的
        # 每个节点应该发的节点是exact one的（未来可以试下有多个node，自己算一个传输快的）
        Q = defaultdict(int)

        max_response_time = 0
        DATA_SIZE = 5
        for own_node in NODES:
            for read_node in NODES:
                translate_speed = NETWORK_SPEED[(own_node, read_node)]

                queue_size = QUEUE_SIZE_PER_NODE[read_node]
                avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
                response_time = queue_size * avg_execution_time +  DATA_SIZE / translate_speed

                if response_time > max_response_time:
                    max_response_time = response_time

        # # 每次只能有一个节点
        exact_one_lagrange = 1  #要是太小了就没办法变成硬constraint了，以后要自适应下
        for time in TIME_FRAMES:
            for own_node in NODES:  #实际持有qubit的node
                qubits = [getNodeQubit(own_node, read_node, time,) for read_node in NODES]
                makeExactOne(Q, qubits, exact_one_lagrange)

        # # 必须得在最快被响应
        # # 假设数据大小固定 (queue size * average execution time + data size / network size)

        time_cost_lagrange = 0.5 / max_response_time
        for time in TIME_FRAMES:
            for own_node in NODES:  #实际持有qubit的node
                for read_node in NODES:
                    translate_speed = NETWORK_SPEED[(own_node, read_node)]

                    queue_size = QUEUE_SIZE_PER_NODE[read_node]
                    avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
                    response_time = queue_size * avg_execution_time +  DATA_SIZE / translate_speed

                    qubit = getNodeQubit(own_node, read_node, time)
                    Q[(qubit, qubit)] += response_time * time_cost_lagrange

        # 互相之间尽量不要发到一起，如果一样，总时长会加avg_execution_time
        for time in TIME_FRAMES:
            for read_node in NODES:
                for own_node1, own_node2 in gridSweep(NODES):
                    qubit1 = getNodeQubit(own_node1, read_node, time)
                    qubit2 = getNodeQubit(own_node2, read_node, time)
                    avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
                    Q[(qubit1, qubit2)] += avg_execution_time * time_cost_lagrange

        # # 如果一个node被add_queue了（也就是为1），后面的对应avg_execution_time的node被选中的概率都会被depress
        for time in TIME_FRAMES:
            for own_node in NODES:
                for read_node in NODES:
                    now_read_qubit = getNodeQubit(own_node, read_node, time)

                    avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
                    effect_times = TIME_FRAMES[time+1: time+int(avg_execution_time)]  #被前面影响的时间
                    
                    for effect_own_node in NODES:
                        for effect_time in effect_times:
                            effect_read_qubit = getNodeQubit(effect_own_node, read_node, effect_time)
                            Q[(now_read_qubit, effect_read_qubit)] += avg_execution_time * time_cost_lagrange   # now_read_qubit * avg_execution_time * later_read_qubit  如果前一个已经被分配了一个queue task那后一个的等待时间就会加一个avg_execution_time

        # 需不需要加入一个goal，今早算完：不需要，因为不知道有多少任务（越早开始是否表示越早结束）

        # 假设当前定了，状态会变成这个定的


        # csp = dwavebinarycsp.ConstraintSatisfactionProblem(dwavebinarycsp.BINARY)

        # bqm = dwavebinarycsp.stitch(csp)

        # https://link.springer.com/article/10.1007/s10951-019-00629-3
        # online schedule algorithm

        bqm = dimod.BQM.from_qubo(Q) # (A, B) and (B, A) will be mergeed automatically
        self.bqm = bqm
        self.Q = Q
        return bqm


    def scheduleAll(self, tasks):
        if self.former_hearbeats is None or self.former_hearbeats is not self.heartbeats: # is 是判断指针的 {} is {} 是false
            self.refreshBQM()

        bqm = self.bqm
        results = sampler.sample(bqm, num_reads=1)
        sample, energy, num_occurrences = list(results.data(['sample', 'energy', 'num_occurrences']))[0]
        sample.update(self.fixed_qubits)

        frame2node2assign_node = praise2frame_qubit(sample, len(self.nodes), self.heartbeat_rate)

        targets = []
        for task in tasks:
            start_node_index = task.start_node_index
            target_node_index = frame2node2assign_node[self.inner_frame_index][start_node_index]
            targets.append((target_node_index, task))

            for read_node in self.nodes:
                read_node_index = read_node.index
                qubit = getNodeQubit(start_node_index, read_node_index, self.inner_frame_index)
                if read_node_index == target_node_index:
                    bqm.fix_variable(qubit, 1)
                    self.fixed_qubits[qubit] = 1
                else:
                    bqm.fix_variable(qubit, 0)
                    self.fixed_qubits[qubit] = 0


            assert target_node_index is not None  #可能会有exact one设置的太低了出现全0或者多个1
        
        self.inner_frame_index += 1
        return targets

if __name__ == '__main__':
    system = QuantumSystemNoResource()
    system.createNodes(5)
    for time_frame in system.time_frames:
        print(time_frame)
        system.step()
        pass

    analysis(system)
    pass